'''
2-4 训练网络
训练集样本总数70，猫35，狗35
'''
import numpy as np
import torch
import torch.nn as nn
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from net_model import ICNET  #导入网络模型，ICNET创建的网络的类名，net_model是储存网络 .py 的文件名
net = ICNET()  # 后面的训练程序可以直接使用 net 来作为网络的接口

x = np.load(file="cat_train_set.npy") / 255   #载入训练集并进行简单归一化处理
x = torch.tensor(x).type(torch.FloatTensor)
y1 = torch.zeros(35)
y2 = torch.ones(35)
y = torch.cat((y1,y2)).type(torch.LongTensor)
optimizer = torch.optim.SGD(net.parameters(), lr=0.02*0.99)   #设置优化器
loss_func = nn.CrossEntropyLoss()      #设置损失函数

samplenum = 70  # 训练集总数
minibatch = 14   # 小批次样本大小

#使用小批次训练的方法进行训练，需要用到 epoch 和 iterations
w_HR = 128       # 样本尺寸
x0 = np.zeros(minibatch * 3 * w_HR * w_HR)
x0 = np.reshape(x0,(minibatch,3,w_HR,w_HR))    #创建小批次空白样本
y0 = np.zeros(minibatch)                       #创建小批次空白标签
x0 = torch.tensor(x0).type(torch.FloatTensor)
y0 = torch.tensor(y0).type(torch.LongTensor)


plt.ion()  #开启交互模式
x_plt = [0] # x坐标
y_plt = [0] # y坐标
for epoch in range(100):                                  # epoch 循环
    for iterations in range(int(samplenum / minibatch)):  # iterations 循环
        k = 0
        for i in range(iterations * minibatch, iterations * minibatch + minibatch):   # 部分样本赋值给 x0 的循环
            x0[k, 0, :, :] = x[i, 0, :, :]
            x0[k, 1, :, :] = x[i, 1, :, :]
            x0[k, 2, :, :] = x[i, 2, :, :]
            y0[k] = y[i]
            k = k + 1
            out = net(x0)   # 实际输出
            loss = loss_func(out,y0)   #实际输出与期望输出传入损失函数
            optimizer.zero_grad()      #清除梯度
            loss.backward()            #误差反向传播
            optimizer.step()           #开始优化

    if epoch % 5 == 0:
        plt.cla()     # 清除上一次绘图
        plt.xlim((0,100))
        plt.xlabel('epoch')
        plt.ylim((0,1))
        plt.ylabel('loss')
        x_plt.append(epoch)  #增加 x 坐标
        y_plt.append(loss.data) #增加 y 坐标
        plt.plot(x_plt,y_plt,c='r',marker='x')         #绘制折线图
        print(f'epoch:{epoch},loss:{loss}')       #输出中间过程
        plt.pause(0.1)       # 停留显示

plt.ioff()  # 关闭交互模式
plt.show()  # 显示最后一张图

torch.save(net,'net.pkl')  #保存网络









